#include <stdlib.h>
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include <opencv2/highgui/highgui_c.h>
#include <thread>
#include <mutex>
#include <unistd.h>
#include <matio.h>
#include <robot/kuka/carmclient.h>
#include <robot/sdh/chandclient.h>
#include <buola/utility/random.h>
#include <buola/app/ccmdline.h>

static std::vector<double> sSensorData[2];
static size_t sSensorRows[2];
static size_t sSensorCols[2];
static bool sHaveSensorData=false;

std::mutex sMutex;

using namespace cv;
using namespace buola;

void sdh_thread();
void hand_thread();

static const double sHeights[]={0.02,0.07,0.12};
static const double sForces[]={1.0};

static const double sHandOpen[]={90_deg,-90_deg,0_deg,-90_deg,-90_deg,-90_deg,0_deg};
static const double sHandClose[]={90_deg,0_deg,45_deg,-90_deg,-90_deg,0_deg,45_deg};

static mat::CVecNd<7> sHandPose;
static robot::sdh::CHandClient *sHandClient;

struct CFrame
{
    int mIndex;
    mat::CVecNd<7> mHandPose;
    std::vector<double> mArmPose;
    int64_t mTimeStamp;
    double mCurrent;
    Mat mDisparity;
    Mat mImage;
    std::vector<double> mSensorData[2];
};

void save_frame(const std::string &pFolder,CFrame &pFrame)
{
    mkdir(pFolder.c_str(),0755);

    // save to file
    std::string lFileName=pFolder+"/image_"+to_string(pFrame.mIndex,5,'0')+".png";
    imwrite(lFileName,pFrame.mImage);
    lFileName=pFolder+"/disp_"+lNumber+".dat";
    FILE *fout = fopen(lFileName.c_str(),"wb");
    fwrite(pFrame.mDisparity.data, sizeof(float), 640*480, fout);
    fclose(fout);

    lFileName=pFolder+"/sensors_"+lNumber+".mat";
    mat_t *lMatFile=Mat_Open(lFileName.c_str(),MAT_ACC_RDWR);
    for(int i=0;i<2;i++)
    {
        //actually variable name
        std::string lVarName="array_"+to_string(i+1);
        size_t lDims[2]={sSensorRows[i],sSensorCols[i]};
        matvar_t *lMatVar=Mat_VarCreate(lVarName.c_str(),MAT_C_DOUBLE,MAT_T_DOUBLE,2,lDims,pFrame.mSensorData[i].data(),0);
        Mat_VarWrite(lMatFile,lMatVar,MAT_COMPRESSION_NONE);
        Mat_VarFree(lMatVar);
    }
    Mat_Close(lMatFile);
    
    lFileName=pFolder+"/info_"+lNumber+".txt";
    io::CTextWriter lInfoFile(io::file(lFileName));
    lInfoFile << pFrame.mIndex << "\n" << pFrame.mTimeStamp << "\n";
    for(int i=0;i<7;i++)
    {
        if(i!=0) lInfoFile << " ";
        lInfoFile << pFrame.mHandPose[i];
    }
    lInfoFile << "\n";
    for(int i=0;i<6;i++)
    {
        if(i!=0) lInfoFile << " ";
        lInfoFile << pFrame.mArmPose[i];
    }
    lInfoFile << "\n";
    lInfoFile << pFrame.mCurrent << "\n";
}

int main( int argc, char* argv[] )
{
    buola_init(argc,argv);
    
    if(cmd_line().ArgCount()!=1)
    {
        msg_info() << "usage: grabsensors folder\n";
        return 1;
    }
    
    std::thread lSensorThread(sdh_thread);

    VideoCapture capture;
    capture.open(CV_CAP_OPENNI);

    if (!capture.isOpened())
        msg_info() << "capture open failed\n";
    
    capture.set(CV_CAP_PROP_OPENNI_REGISTRATION,1);

    while(!sHaveSensorData) usleep(10000);
    
    robot::kuka::CArmClient lArm;
    lArm.Init();
    robot::sdh::CHandClient lHand;
    lHand.Init();
    sHandClient=&lHand;
    std::thread lHandThread(hand_thread);
    
    sleep(1_s);
    int lFrameIndex=1;

    lHand.SetGoal(sHandOpen);
    lHand.WaitForStop();
    
    for(int lIter=0;lIter<50;lIter++)
    {
        for(double lHeight : sHeights)
        {
            msg_info() << "moving arm\n";
            start_timer();
            mat::CVec3d lArmPos(0,-0.7+random::get<double>(-0.02,0.02),
                                lHeight+random::get<double>(-0.02,0.02));
            msg_info() << "moving arm to " << lArmPos << "\n";
            lArm.SetGoal(lArmPos);
            lArm.WaitForStop();
            end_timer();
            
            for(double lForce : sForces)
            {
                lFrameIndex=1;
                msg_info() << "closing hand\n";
                start_timer();
                lHand.SetCurrent(lForce);
                lHand.SetGoal(sHandClose);
                int lStopCount=0;
                std::vector<CFrame*> lFrames;
                while(lStopCount<10)
                {
                    if(!lHand.IsMoving()) lStopCount++;
                    
                    CFrame *lFrame=new CFrame;
                    Mat lDepthMap,lShowDepth,lImage,lDisparity;
                    
                    capture.grab();
                    capture.retrieve( lDepthMap, CV_CAP_OPENNI_DEPTH_MAP );
                    capture.retrieve( lImage, CV_CAP_OPENNI_BGR_IMAGE );
                    capture.retrieve( lDisparity, CV_CAP_OPENNI_DISPARITY_MAP_32F );
                    lFrame->mImage=lImage.clone();
                    lFrame->mDisparity=lDisparity.clone();

                    lDepthMap.convertTo(lShowDepth, CV_8UC1, 255.0/5000.0);
                    
                    imshow("bgr",lImage);
                    imshow("depth",lShowDepth);
                    lArm.GetPose(lFrame->mArmPose);
                    lFrame->mIndex=lFrameIndex++;
                    lFrame->mCurrent=lForce;
                    lFrame->mTimeStamp=CCalendarClock::now().time_since_epoch().count();
                    sMutex.lock();
                    lFrame->mSensorData[0]=sSensorData[0];
                    lFrame->mSensorData[1]=sSensorData[1];
                    lFrame->mHandPose=sHandPose;
                    sMutex.unlock();
                    lFrames.push_back(lFrame);
                    waitKey(1);
                }
                mkdir(cmd_line().Arg(0).c_str(),0755);
                std::string lFolder=cmd_line().Arg(0)+"/set_"+to_string(lHeight)+"_"+to_string(lForce)+"_"+to_string(lIter);
                bool lFoundContact=false;
                lFrameIndex=1;
                for(CFrame *lFrame : lFrames)
                {
                    if(lFrameIndex<10)
                    {
                        lFrame->mIndex=lFrameIndex++;
                        save_frame(lFolder,*lFrame);
                        delete lFrame;
                        
                        continue;
                    }
                    
                    if(!lFoundContact)
                    {
                        double lForce1=std::accumulate(lFrame->mSensorData[0].begin(),lFrame->mSensorData[0].end(),0.0);
                        double lForce2=std::accumulate(lFrame->mSensorData[1].begin(),lFrame->mSensorData[1].end(),0.0);
                        
                        if(lForce1+lForce2>0)
                            lFoundContact=true;
                        else
                            continue;
                    }
                    
                    lFrame->mIndex=lFrameIndex++;
                    save_frame(lFolder,*lFrame);
                    delete lFrame;
                }
                lHand.WaitForStop();
                end_timer();
                msg_info() << "opening hand\n";
                start_timer();
                lHand.SetGoal(sHandOpen);
                lHand.WaitForStop();
                end_timer();
            }
        }
    }

    return buola_finish();
}

#include <sdh/sdh.h>
#include <sdh/dsa.h>
#include <sdh/basisdef.h>
#include <sdh/util.h>

//----------------------------------------------------------------------
// Defines, enums, unions, structs,
//----------------------------------------------------------------------

USING_NAMESPACE_SDH

int sIndices[]={1,5};

void sdh_thread()
{
    try
    {
        cDSA ts = cDSA(0, 0);

/*
        // Set sensitivities if requested
        for ( i=0; i<6; i++ )
        {
            if ( options.sensitivity[i] < 0.0 )
                continue;
            ts.SetMatrixSensitivity(i, options.sensitivity[i], false, options.reset_to_default, options.persistent );
        }

        // Set thresholds if requested
        for ( i=0; i<6; i++ )
        {
            if ( options.threshold[i] > 4095 )
                continue;
            ts.SetMatrixThreshold(i, options.threshold[i], false, options.reset_to_default, options.persistent );
        }
*/

        //---------------------
        // Read and show actual tactile sensor data ("full frame")

        // Make remote tactile sensor controller send data automatically as fast as possible (DSA push-mode):
        ts.SetFramerate( 1, true);


            //-----------
            // start periodic or one time processing of full frames if requested:
            while(true)
            {
                //-----------
                try
                {
                    ts.UpdateFrame();
                    std::vector<double> lSensorData[2];
                    
                    for(int i=0;i<2;i++)
                    {
                        int m=sIndices[i];
                    
                        lSensorData[i].clear();
                        
                        sSensorRows[i]=ts.GetMatrixInfo(m).cells_y;
                        sSensorCols[i]=ts.GetMatrixInfo(m).cells_x;
                    
                        for(int x=0;x<ts.GetMatrixInfo(m).cells_x;x++)
                        {
                            for(int y=0;y<ts.GetMatrixInfo(m).cells_y;y++)
                            {
                                lSensorData[i].push_back(ts.GetTexel(m,x,y));
                            }
                        }
                    }
                    
                    sMutex.lock();
                    sSensorData[0]=std::move(lSensorData[0]);
                    sSensorData[1]=std::move(lSensorData[1]);
                    sHaveSensorData=true;
                    sMutex.unlock();
                }
                catch ( cDSAException* e)
                {
                    msg_info() << "Caught and ignored cDSAException: " << e->what() << "\n";
                    delete e;
                }
            }

        ts.Close();
    }
    catch ( cSDHLibraryException* e )
    {
        msg_info() << "\ndemo-dsa main(): Caught exception from SDHLibrary: " << e->what() << ". Giving up!\n";
        delete e;
    }
    catch (...)
    {
        msg_info() << "\ncaught unknown exception, giving up\n";
    }
}

void hand_thread()
{
    while(true)
    {
        mat::CVecNd<7> lPose;
        sHandClient->GetPose(lPose);
        sMutex.lock();
        sHandPose=lPose;
        sMutex.unlock();
    }
}
